from numpy.fft import fft
import numpy as np

#Este material foi adaptado do projeto Pyftd sobre a licenca MIT
# fonte :http://code.google.com/p/pytfd/

def subset(x, i, N ):
    assert N <= len(x)
    if i < N//2:
        x_subset = x[:i + N//2]
        x_subset = np.concatenate([np.zeros(N//2 - i), x_subset])
    else:
        x_subset = x[i - N//2 : i + N//2]
        if len(x_subset) < N:
            x_subset = np.concatenate([x_subset, np.zeros(N - len(x_subset))])
    return x_subset

def zeropad(w, N):
    T = len(w)
    w_zeros = np.zeros(N//2 - T//2)
    w = np.concatenate([w_zeros, w, w_zeros])
    if len(w) == N + 1:
        w = w[:-1]
    return w

#X eh um vetor
# w uma janela
# L o overlap
def stft(x, w, L=None):
    N = len(x)
    if L is None:
        L = N
    w = zeropad(w, N)
    X_stft = []
    points = range(0, N, N//L)
    for i in points:
        x_subset = subset(x, i, N)
        fft_subset = fft(x_subset * w)
        X_stft.append(fft_subset)
    X_stft = np.array(X_stft).transpose()
    return X_stft
